package com.fjj.cs.util;

import lombok.Data;

import java.util.*;

/**
 * 定义Trie类
 */
@Data
public class Trie {
    // 根节点
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    // 插入单词
    public void insert(String word) {
        // 从根节点开始
        TrieNode current = root;
        // 遍历单词中的每个字符
        for (char ch : word.toCharArray()) {
            // 如果当前字符没有对应的子节点，则创建一个新的节点
            current.children.putIfAbsent(ch, new TrieNode());
            // 移动到子节点
            current = current.children.get(ch);
        }
        // 最后一个字符的节点标记为单词的结尾
        current.isEndOfWord = true;
    }

    // 搜索单词
    public boolean search(String word) {
        // 调用searchNode方法查找指定单词对应的最后一个节点
        TrieNode node = searchNode(word);
        // 如果节点存在且是单词结尾，则返回true，否则返回false
        return node != null && node.isEndOfWord;
    }

    // 检查是否有以给定前缀开头的单词
    public boolean startsWith(String prefix) {
        // 调用searchNode方法查找指定前缀对应的最后一个节点
        return searchNode(prefix) != null;
    }

    // 查找所有以给定前缀开头的单词
    public List<String> findWordsWithPrefix(String prefix) {
        List<String> result = new ArrayList<>();
        // 查找前缀对应的节点
        TrieNode node = searchNode(prefix);
        if (node != null) {
            // 如果前缀节点存在，收集所有以该前缀开头的单词
            collectWords(node, new StringBuilder(prefix), result);
        }
        // 返回结果列表
        return result;
    }

    // 查找给定字符串（单词或前缀）对应的最后一个节点
    private TrieNode searchNode(String str) {
        TrieNode current = root;
        // 遍历字符串中的每个字符
        for (char ch : str.toCharArray()) {
            // 如果当前字符没有对应的子节点，返回null
            if (!current.children.containsKey(ch)) {
                return null;
            }
            // 移动到子节点
            current = current.children.get(ch);
        }
        // 返回最后一个节点
        return current;
    }

    // 递归地收集所有以当前节点为起点的单词
    private void collectWords(TrieNode node, StringBuilder prefix, List<String> result) {
        // 如果当前节点是单词的结尾，添加当前的前缀到结果列表
        if (node.isEndOfWord) {
            result.add(prefix.toString());
        }
        // 遍历当前节点的所有子节点，递归收集以当前节点为起点的单词
        for (Map.Entry<Character, TrieNode> entry : node.children.entrySet()) {
            // 将当前字符添加到前缀中
            prefix.append(entry.getKey());
            // 递归收集子节点的单词
            collectWords(entry.getValue(), prefix, result);
            // 回溯，删除当前字符，继续遍历其他子节点
            prefix.deleteCharAt(prefix.length() - 1);
        }
    }
}